//
// Created by wzc on 2023/2/6.
//

// virtual_depot -> drone_pos -> ftr-x -> ... -> ftr-x -> virtual_depot
// srv.request.prob = 1 求解tsp或者mtsp问题，这个tsp问题是有min size和max size要求的
// srv.request.prob = 2 求解tsp后者mtsp问题，没有min size和max size要求
// srv.request.prob = 3 求解acvrp问题

#include <ros/ros.h>
#include <ros/package.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <visualization_msgs/Marker.h>
#include <Eigen/Eigen>
#include <iostream>
#include <fstream>

#include <lkh_mtsp_solver/SolveMTSP.h>

std::string mtsp_dir = ros::package::getPath("lkh_mtsp_solver") + "/resource2";

ros::Publisher vis_pub_;

std::vector<Eigen::Vector3d> drone_pos;
std::vector<Eigen::Vector3d> ftr_pos;
std::vector<Eigen::Vector3d> virtual_depot;

ros::ServiceClient tsp_client;

bool solved;

void drawPoint(const int &id, const std::string &ns, const Eigen::Vector3d &pos,
               const Eigen::Vector3d &scale, const Eigen::Vector4d &color) {
  visualization_msgs::Marker marker;
  marker.header.frame_id = "world";
  marker.header.stamp = ros::Time::now();
  marker.ns = ns;
  marker.id = id;
  marker.action = visualization_msgs::Marker::ADD;
  marker.type = visualization_msgs::Marker::SPHERE;
  marker.pose.position.x = pos[0];
  marker.pose.position.y = pos[1];
  marker.pose.position.z = pos[2];
  marker.pose.orientation.x = 0.0;
  marker.pose.orientation.y = 0.0;
  marker.pose.orientation.z = 0.0;
  marker.pose.orientation.w = 1.0;
  marker.scale.x = scale[0];
  marker.scale.y = scale[1];
  marker.scale.z = scale[2];
  marker.color.r = color[0];
  marker.color.g = color[1];
  marker.color.b = color[2];
  marker.color.a = color[3];

  vis_pub_.publish(marker);
  ros::Duration(0.0005).sleep();

  marker.ns = ns + "_id";
  marker.id = id;
  marker.action = visualization_msgs::Marker::ADD;
  marker.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
  marker.pose.position.x = pos[0] + 0.2;
  marker.pose.position.y = pos[1] + 0.2;
  marker.pose.position.z = pos[2] + 0.4;
  marker.text = ns + "_" + std::to_string(id);
  marker.scale.x = 0.5;
  marker.scale.y = 0.5;
  marker.scale.z = 0.5;
  vis_pub_.publish(marker);
  ros::Duration(0.0005).sleep();
}

void drawResult(const std::vector<int> &result, const Eigen::Vector3d &pos, const int &id) {
  visualization_msgs::Marker marker;
  marker.header.frame_id = "world";
  marker.header.stamp = ros::Time::now();
  marker.ns = std::to_string(id) +"_result_start";
  marker.id = 0;
  marker.action = visualization_msgs::Marker::ADD;
  marker.type = visualization_msgs::Marker::LINE_STRIP;
  marker.pose.position.x = 0.0;
  marker.pose.position.y = 0.0;
  marker.pose.position.z = 0.0;
  marker.pose.orientation.x = 0.0;
  marker.pose.orientation.y = 0.0;
  marker.pose.orientation.z = 0.0;
  marker.pose.orientation.w = 1.0;
  marker.scale.x = 0.1;
  marker.scale.y = 0.0;
  marker.scale.z = 0.0;
  marker.color.r = 0.0;
  marker.color.g = 0.0;
  marker.color.b = 0.0;
  marker.color.a = 1.0;

  geometry_msgs::Point pt;
  pt.x = pos(0);
  pt.y = pos(1);
  pt.z = pos(2);
  marker.points.push_back(pt);

  pt.x = ftr_pos[result[0]](0);
  pt.y = ftr_pos[result[0]](1);
  pt.z = ftr_pos[result[0]](2);
  marker.points.push_back(pt);

  vis_pub_.publish(marker);
  ros::Duration(0.005).sleep();

  for(int i = 0; i < result.size()-1; i++) {
    visualization_msgs::Marker marker;
    marker.header.frame_id = "world";
    marker.header.stamp = ros::Time::now();
    marker.ns = std::to_string(id) +"_result";
    marker.id = i;
    marker.action = visualization_msgs::Marker::ADD;
    marker.type = visualization_msgs::Marker::LINE_STRIP;
    marker.pose.position.x = 0.0;
    marker.pose.position.y = 0.0;
    marker.pose.position.z = 0.0;
    marker.pose.orientation.x = 0.0;
    marker.pose.orientation.y = 0.0;
    marker.pose.orientation.z = 0.0;
    marker.pose.orientation.w = 1.0;
    marker.scale.x = 0.1;
    marker.scale.y = 0.0;
    marker.scale.z = 0.0;
    marker.color.r = 0.0;
    marker.color.g = 0.0;
    marker.color.b = 0.0;
    marker.color.a = 1.0;

    geometry_msgs::Point pt;
    pt.x = ftr_pos[result[i]](0);
    pt.y = ftr_pos[result[i]](1);
    pt.z = ftr_pos[result[i]](2);
    marker.points.push_back(pt);

    pt.x = ftr_pos[result[i+1]](0);
    pt.y = ftr_pos[result[i+1]](1);
    pt.z = ftr_pos[result[i+1]](2);
    marker.points.push_back(pt);

    vis_pub_.publish(marker);
    ros::Duration(0.005).sleep();
  }
}

void init() {
  solved = false;

  std::cout << "[Init]init done, test mTSP!" << std::endl;
  std::cout << "[Init]use /move_base_simple/goal to give waypoints" << std::endl;
  std::cout << "[Init]use /initialpose to give start" << std::endl;
}

void solveTSP() {
  std::cout << "[solveTSP]start build cost matrix" << std::endl;
  int drone_num = drone_pos.size();
  int ftr_num = ftr_pos.size();
  int dim = 1 + drone_num + ftr_num;
  Eigen::MatrixXd mat = Eigen::MatrixXd::Zero(dim, dim);

  // virtual depot to drones
  for(int i = 0; i < drone_num; i++ ) {
    mat(0, 1 + i) = -1000;
    mat(1 + i, 0) = 1000;
  }

  // virtual depot to frontiers
  for (int i = 0; i < ftr_num; i++) {
    mat(0, 1 + drone_num + i) = 1000;
    mat(1 + drone_num + i, 0) = 0;
  }

  // cost between drones
  for (int i = 0; i < drone_num; i++) {
    for (int j = 0; j < drone_num; j++) {
      mat(1 + i, 1 + j) = 10000;
    }
  }

  // cost form drone to frontiers
  for (int i = 0; i < drone_num; i++) {
    for (int j = 0; j < ftr_num; j++) {
      mat(1 + i, 1 + drone_num + j) = (drone_pos[i] - ftr_pos[j]).norm();
      mat(1 + drone_num + j, 1 + i) = 0;
    }
  }

  for (int i = 0; i < ftr_num; i++) {
    for (int j = 0; j < ftr_num; j++) {
      mat(1 + drone_num + i, 1 + drone_num + j) = (ftr_pos[i] - ftr_pos[j]).norm();
    }
  }
  // diag
  for (int i = 0; i < dim; i++) {
    mat(i, i) = 1000;
  }

  // vis cost matrix
  std::cout << "[solveTSP]=====cost matrix build done=====" << std::endl;
  for (int i = 0; i < dim; i++) {
    for (int j = 0; j < dim; j++) {
      std::cout << mat(i, j) << ", ";
    }
    std::cout << std::endl;
  }
  std::cout << "========================================" << std::endl;

  std::cout << "[solveTSP]=====write cost matrix ot atsp2.atsp======" << std::endl;
  std::ofstream file(mtsp_dir + "/amtsp2_1.atsp");
  file << "NAME : amtsp\n";
  file << "TYPE : ATSP\n";
  file << "DIMENSION : " + std::to_string(dim) + "\n";
  file << "EDGE_WEIGHT_TYPE : EXPLICIT\n";
  file << "EDGE_WEIGHT_FORMAT : FULL_MATRIX\n";
  file << "EDGE_WEIGHT_SECTION\n";
  for (int i = 0; i < dim; ++i) {
    for (int j = 0; j < dim; ++j) {
      int int_cost = 100 * mat(i, j);
      file << int_cost << " ";
    }
    file << "\n";
  }
  file.close();

  std::cout << "[solveTSP]=====write cost matrix ot atsp2.par======" << std::endl;
  // Create par file
  file.open(mtsp_dir + "/amtsp2_1.par");
  file << "SPECIAL\n";
  file << "PROBLEM_FILE = " + mtsp_dir + "/amtsp2_1.atsp\n";
  file << "SALESMEN = " << std::to_string(drone_num) << "\n";
  file << "MTSP_OBJECTIVE = MINSUM\n";
  file << "MTSP_MIN_SIZE = " << std::to_string(1) << "\n"; // 可以删除，这个是控制经过点的min数量的参数
  file << "MTSP_MAX_SIZE = " << std::to_string(10) << "\n"; // 可以删除，这个是控制经过点的max数量的参数
  file << "RUNS = 1\n";
  file << "TRACE_LEVEL = 0\n";
  file << "TOUR_FILE = " + mtsp_dir + "/amtsp2_1.tour\n";
  file.close();

  ros::Time t1 = ros::Time::now();
  lkh_mtsp_solver::SolveMTSP srv;
  srv.request.prob = 2;
  if (!tsp_client.call(srv)) {
    ROS_ERROR_STREAM("Fail to solve ATSP.");
    return;
  }
  double solve_t = (ros::Time::now() - t1).toSec();

  ROS_WARN_STREAM("[solveTSP]tsp solve t " << solve_t);

  std::ifstream fin(mtsp_dir + "/amtsp2_1.tour");
  std::string res;
  std::vector<int> ids;
  while (getline(fin, res)) {
    if (res.compare("TOUR_SECTION") == 0) break;
  }
  while (getline(fin, res)) {
    int id = stoi(res);
    ids.push_back(id - 1);
    if (id == -1) break;
  }
  fin.close();

  // Parse the m-tour
  std::vector<std::vector<int>> tours;
  std::vector<int> tour;
  for (auto id : ids) {
    if (id > 0 && id <= drone_num) {
      tour.clear();
      tour.push_back(id);
    } else if (id >= dim|| id <= 0) {
      tours.push_back(tour);
    } else {
      tour.push_back(id);
    }
  }

  std::cout << "[solveTSP] tours size is " << tours.size() << std::endl;
  for(int j = 0; j < tours.size(); j++) {
    std::cout << "[solveTSP] tours " << j << " ";
    for (int i = 0; i < tours[j].size(); i++) {
      if (i==tours[j].size()-1) {
        std::cout << tours[j][i];
      } else {
        std::cout << tours[j][i] << "->";
      }
    }
    std::cout << std::endl;
  }

  for(int k = 0; k < tours.size(); k++) {
    if (tours[k].empty()) continue;
    std::vector<int> result;
    result.insert(result.end(), tours[k].begin() + 1, tours[k].end());

    for (auto& id : result) {
      id -= 1 + drone_num;
    }

    std::cout << "[solveTSP] result " << k <<  " is: ";
    for (int i = 0; i < result.size(); i++) {
      if (i==result.size()-1) {
        std::cout << result[i];
      } else {
        std::cout << result[i] << "->";
      }
    }
    std::cout << std::endl;
    drawResult(result, drone_pos[tours[k][0]-1], k);
  }
}

void GoalCallback(const geometry_msgs::PoseStampedPtr &msg) {
  if (solved) {
    ROS_ERROR_STREAM("solved, please restart");
    return;
  }

  Eigen::Vector3d ftr;
  ftr << msg->pose.position.x, msg->pose.position.y, 1.0;
  ftr_pos.push_back(ftr);
  std::cout << "[Goal]Get num " << ftr_pos.size() << " ftr." << std::endl;


  drawPoint(ftr_pos.size()-1, "way_pt", ftr, Eigen::Vector3d(0.2, 0.2, 0.2), Eigen::Vector4d(0, 0, 1, 1));
}

void triggerCallback(const geometry_msgs::PoseWithCovarianceStampedConstPtr &msg) {
  const int total_drone_num = 3;

  if (solved) {
    ROS_ERROR_STREAM("solved, please restart");
    return;
  }

  if (ftr_pos.size() < 3) {
    ROS_ERROR_STREAM("ftr num to small!");
    return;
  }

  Eigen::Vector3d drone;
  drone << msg->pose.pose.position.x, msg->pose.pose.position.y, 1.0;
  drone_pos.push_back(drone);
  std::cout << "[Init]drone num " << drone_pos.size() << "/" << total_drone_num << std::endl;

  drawPoint(drone_pos.size()-1, "start", drone_pos[drone_pos.size()-1], Eigen::Vector3d(0.4, 0.4, 0.4), Eigen::Vector4d(1, 0, 0, 1));

  if (drone_pos.size() == total_drone_num) {
    solveTSP();
    solved = true;
  }
}


int main(int argc, char** argv)
{
  ros::init(argc,argv,"test_tsp");
  ros::NodeHandle nh("~");

  ros::Duration(4.0).sleep();

  tsp_client = nh.serviceClient<lkh_mtsp_solver::SolveMTSP>("/solve_tsp_1", true);
  vis_pub_ = nh.advertise<visualization_msgs::Marker>("/tsp_vis", 1000);

  ros::Subscriber goal_sub = nh.subscribe("/move_base_simple/goal", 10, GoalCallback);
  ros::Subscriber trigger_sub = nh.subscribe("/initialpose", 10, triggerCallback);

  ros::Duration(1.0).sleep();

  init();

  ros::spin();
}